extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/log.h>
}
#include <fstream>
#include <iostream>
using namespace std;
#ifdef main
#undef main
#endif // !main
int main(int argc, char *argv[]) {
  string filename = "400_300_25";
  AVCodecID codec_id = AV_CODEC_ID_H264;
  if (argc > 1) {
    string codec = argv[1];
    if (codec == "h265" || codec == "hevc") {
      codec_id = AV_CODEC_ID_H265;
    } else {
    }
  }
  if (codec_id == AV_CODEC_ID_H264) {
    filename += ".h264";
  } else if (codec_id == AV_CODEC_ID_H265) {
    filename += ".h265";
  }
  ofstream ofs;
  ofs.open(filename, ios::binary);
  // 设置自定义日志回调函数
  av_log_set_level(AV_LOG_DEBUG);
  AVFrame *frame = nullptr;
  AVPacket *packet = nullptr;
  // 1 查找编码器
  AVCodec *codec = avcodec_find_encoder(codec_id);
  if (!codec) {
    cout << "Cannot find encoder for H264" << endl;
    // 尝试查找H265编码器
    codec = avcodec_find_encoder(AV_CODEC_ID_HEVC);
    if (!codec) {
      cout << "Cannot find encoder for HEVC either" << endl;
      return -1;
    } else {
      cout << "Found HEVC encoder" << endl;
    }
  }

  // 2 创建编码器上下文
  AVCodecContext *c = avcodec_alloc_context3(codec);
  if (!c) {
    cout << "Cannot alloc codec context" << endl;
    return -1;
  }

  // 3 设置编码参数
  c->width = 640;
  c->height = 480;
  c->time_base = AVRational{1, 25}; // 一帧的时间为1/25秒
  c->pix_fmt = AV_PIX_FMT_YUV420P;    // 指定源视频的像素格式
  c->thread_count = 8;

  auto ret = avcodec_open2(c, codec, NULL);
  if (ret != 0) {
    char buf[1024] = {0};
    av_strerror(ret, buf, sizeof(buf));
    cout << "Cannot open codec: " << buf << endl;
    goto err;
  }
  cout << "Codec opened successfully" << endl;
  frame = av_frame_alloc();
  if (!frame) {
    char buf[1024] = {0};
    av_strerror(ret, buf, sizeof(buf));
    cout << "Cannot open codec: " << buf << endl;
    goto err;
  }
  frame->width = c->width;
  frame->height = c->height;
  frame->format = c->pix_fmt;
  frame->linesize[0] = c->width;
  frame->linesize[1] = c->width / 2;
  frame->linesize[2] = c->width / 2;
  ret = av_frame_get_buffer(frame, 0);
  if (ret != 0) {
    char buf[1024] = {0};
    av_strerror(ret, buf, sizeof(buf));
    cout << "Cannot open codec: " << buf << endl;
    goto err;
  }
  packet = av_packet_alloc();
  for (int i = 0; i < 250; ++i) {
    // 生成AVFrame数据
    // Y
    for (int y = 0; y < c->height; y++) {
      for (int x = 0; x < c->width; x++) {
        frame->data[0][y * frame->linesize[0] + x] = x + y + i * 3;
      }
    }
    for (int y = 0; y < c->height / 2; y++) {
      for (int x = 0; x < c->width / 2; x++) {
        frame->data[1][y * frame->linesize[1] + x] = 128 + y + i * 3;
        frame->data[2][y * frame->linesize[2] + x] = 64 + y + i * 5;
      }
    }
    frame->pts = i;
    ret = avcodec_send_frame(c, frame);
    if (ret < 0) {
      break;
    }
    while (ret >= 0) {
      ret = avcodec_receive_packet(c, packet);
      if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
        break;
      }
      if (ret < 0) {
        char buf[1024] = {0};
        av_strerror(ret, buf, sizeof(buf));
        cout << "Cannot open codec: " << buf << endl;
        goto err;
      }
      cout << packet->size << " " << flush;
      ofs.write((char *)packet->data, packet->size);
      ofs.flush();
      av_packet_unref(packet);
    }
  }
err:
  if (c)
    avcodec_free_context(&c);
  if (frame)
    av_frame_free(&frame);
  if (packet)
    av_packet_free(&packet);
  return 0;
}
